#!/usr/bin/env python

""" Example for integrating pyaiengine with netflow """

__author__ = "Luis Campo Giralte"
__copyright__ = "Copyright (C) 2013-2020 by Luis Campo Giralte"
__revision__ = "$Id$"
__version__ = "0.1"

import sys
sys.path.append("../src/")
import pyaiengine
from scapy.all import *

class netflowAdaptor(pyaiengine.DatabaseAdaptor):
    """ This class inheritance of DatabaseAdaptor that contains 
        the following methods:
        - insert, called on the first insertion of the network flow
        - update, called depending on the sample selected.
        - remove, called when the flow is destroy.
    """
    def __init__(self):
        """Setup the fields and the netflow headers needed
        for send later on the network or pcap."""

        self.__seq = 0

        # Build the netflow template
        self.__flowset = NetflowFlowsetV9(templates=[
            NetflowTemplateV9(
                template_fields=[
                    NetflowTemplateFieldV9(fieldType=1),  # IN_BYTES
                    NetflowTemplateFieldV9(fieldType=2, fieldLength=4),  # IN_PKTS
                    NetflowTemplateFieldV9(fieldType=4),                 # PROTOCOL
                    NetflowTemplateFieldV9(fieldType=7),                 # L4_SRC_PORT
                    NetflowTemplateFieldV9(fieldType=8),                 # IPV4_SRC_ADDR
                    NetflowTemplateFieldV9(fieldType=11),                # L4_DST_PORT
                    NetflowTemplateFieldV9(fieldType=12),                # IPV4_DST_ADDR
                    NetflowTemplateFieldV9(fieldType=23, fieldLength=4), # OUT_BYTES
                    NetflowTemplateFieldV9(fieldType=24, fieldLength=4), # OUT_PKTS
                    NetflowTemplateFieldV9(fieldType=96, fieldLength=8), # APPLICATION_NAME
                ],
                templateID=256,
                fieldCount=10
            )],
            flowSetID=0)

        self.__recordClass = GetNetflowRecordV9(self.__flowset)
        self.__packet_buffer = []

        # In case you want to write to a file
        # self.pktdump = PcapWriter("banana.pcap", append=True, sync=True)

    def add_flow_record(self, flow):
        """This function will create a netflow record that
        will be sent later on a UDP datagram."""

        r = self.__recordClass(
            IN_BYTES=(flow.upstream_bytes).to_bytes(4, byteorder='big'),
            IN_PKTS=(flow.upstream_packets).to_bytes(4, byteorder="big"),
            PROTOCOL=flow.protocol,
            L4_SRC_PORT=flow.src_port,
            IPV4_SRC_ADDR=flow.src_ip,
            L4_DST_PORT=flow.dst_port,
            IPV4_DST_ADDR=flow.dst_ip,
            OUT_BYTES=(flow.downstream_bytes).to_bytes(4, byteorder='big'),
            OUT_PKTS=(flow.downstream_packets).to_bytes(4, byteorder='big'),
            APPLICATION_NAME=flow.l7protocol) 

        self.__packet_buffer.append(r)

    def connect(self, connection_str):
        ip = connection_str.split(":")[0]
        port = connection_str.split(":")[1]
        self.__header = Ether()/IP()/UDP()

    def update(self, flow, data):

        self.add_flow_record(flow)

        if len(self.__packet_buffer) > 8:
            dataFS = NetflowDataflowsetV9(
                templateID=256,
                records=self.__packet_buffer
                )
            netflow_header = NetflowHeader()/NetflowHeaderV9(packageSequence=self.__seq)
            pkt = self.__header / netflow_header / self.__flowset / dataFS 

            # Write the netflow on a pcap
            # self.pktdump.write(pkt)
            # or send over the network
            # send(pkt)

            self.__seq += 1
            self.__packet_buffer = []
 
    def insert(self, flow):
        self.add_flow_record(flow)

    def remove(self, flow):
        self.add_flow_record(flow)

if __name__ == '__main__':

    # Load an instance of a Network Stack on Lan Network
    st = pyaiengine.StackLan()

    st.tcp_flows = 327680
    st.udp_flows = 163840
 
    # Create a netflowAdaptor object. 

    db = netflowAdaptor()
    db.connect("192.168.0.1:2055")
 
    st.set_udp_database_adaptor(db, 128)
    st.set_tcp_database_adaptor(db, 128)

    filename = "/home/luis/pcapfiles/ists-12.pcap"
    with pyaiengine.PacketDispatcher(filename) as pd:
        pd.stack = st
        pd.run()

    sys.exit(0)

